from milvus import Milvus, DataType   # milvus 0.11.0 pymilvus 0.3.0
#from milvus import Milvus, IndexType, MetricType, Status # milvus 0.10.5 pymilvus 0.2.15


# class FaceMilvus0105:
#       # milvus 0.10.5 pymilvus 0.2.15
#       def __init__(self, host='localhost', port='19530', collection_name='face_encode',  dim=128, INDEX_FILE_SIZE=32):
#             self.dim=dim
#             self.milvus = Milvus(host, port)
#             self.collection_name = collection_name
#             self.INDEX_FILE_SIZE = INDEX_FILE_SIZE
#             self.collection = self._create_collection(collection_name)
#
#       def _create_collection(self, collection_name):
#             status, ok = self.milvus.has_collection(collection_name)
#             if not ok :
#                   param = {
#                         'collection_name': collection_name,
#                         'dimension': self.dim,
#                         'index_file_size': self.INDEX_FILE_SIZE,  # optional
#                         'metric_type': MetricType.L2  # optional
#                   }
#                   self.milvus.create_collection(param)
#             # _, collections = self.milvus.list_collections()
#             _, collection = self.milvus.get_collection_stats(collection_name)
#             return collection
#
#       def Insert(self, imgID, imgEncoding, collection_name):
#             vectors = [imgEncoding]
#             status, ids = self.milvus.insert(collection_name=collection_name, records=vectors, ids=[imgID])
#             if not status.OK():
#                   print("Insert failed: {}".format(status))
#             self.Flush()
#
#       def GetEntityByID(self, ids, collection_name=''):
#             name = collection_name or self.collection_name
#             status, result_vectors = self.milvus.get_entity_by_id(self.collection_name, ids[:10])
#             if not status.OK():
#                   print("GetEntityByID failed: {}".format(status))
#             return result_vectors
#
#       def CreateIndex(self, index_type=IndexType.IVF_FLAT, index_param_nlist=2048):
#             index_param = {
#                   'nlist': index_param_nlist
#             }
#             self.milvus.create_index(self.collection_name, index_type, index_param)
#
#       def DropIndex(self, index_name):
#             self.milvus.drop_index(self.collection_name, index_name)
#
#       def Search(self, imgEncoding, collection_name, numVectors=10, partition_tags=None, fields=None, timeout=None):
#             search_param = {
#                   "nprobe": 16
#             }
#             param = {
#                   'collection_name': collection_name,
#                   'query_records': [imgEncoding],
#                   'top_k': numVectors,
#                   'params': search_param,
#             }
#             status, entities_list = self.milvus.search(**param)
#             results = []
#             for entities in entities_list:
#                   for topk_film in entities:
#                         result = {}
#                         current_entity = topk_film.entity
#                         # print("- id: {}".format(topk_film.id))
#                         # print("- distance: {}".format(topk_film.distance))
#                         # print("- imgID: {}".format(current_entity.imgID))
#                         result["id"] = topk_film.id
#                         result["distance"] = topk_film.distance
#                         result["imgID"] = current_entity.imgID
#                         results.append(result)
#             return results
#
#       def GetMilvus(self):
#             return self.milvus
#
#       def DeleteByIds(self, ids):
#             # ids = [1,]
#             self.milvus.delete_entity_by_id(self.collection_name, ids)
#             self.Flush()
#
#       def _delete_collection(self, collection_name):
#             self.milvus.drop_collection(collection_name)
#
#       def Delete(self, collection_name):
#             self._delete_collection(collection_name)
#
#
#       def Flush(self):
#             self.milvus.flush([self.collection_name])


class FaceMilvus:
      # just 适配 milvus 0.11.0 pymilvus 0.3.0
      def __init__(self, host='localhost', port='19530', collection_name='face_encode'):
            self.collection = None
            self.collection_name = collection_name
            self.client = Milvus(host=host, port=port)
            self.collection_param = {
                                    "fields": [
                                          {"name": "imgID", "type": DataType.INT64},
                                          {"name": "imgEncoding", "type": DataType.FLOAT_VECTOR, "params": {"dim": 128}}
                                    ],
                                    "segment_row_limit": 100000,
                                    "auto_id": True
                                    }
            self.collection = self.collection if self.collection  else self._create_collection()
      def _create_collection(self):
            print(self.client.list_collections())
            if self.collection_name in self.client.list_collections() :
                  return 
            return self.client.create_collection(self.collection_name, self.collection_param)

      def _delete_collection(self, collection_name):
            self.client.drop_collection(collection_name)
      
      def Delete(self, collection_name):
            self._delete_collection(collection_name)

      def Insert(self, imgID, imgEncoding, collection_name):
            hybrid_entities = [
                                    # {"name": "duration", "values": list_of_int, "type": DataType.INT32},
                                    # {"name": "release_year", "values": list_of_int, "type": DataType.INT64},
                                    # {"name": "embedding", "values": vectors, "type":DataType.FLOAT_VECTOR},
                                    {"name": "imgID", "values": [imgID], "type": DataType.INT64},
                                    {"name": "imgEncoding", "values": [imgEncoding],"type": DataType.FLOAT_VECTOR}
                              ]
            ids = self.client.insert(collection_name, hybrid_entities)
            self.Flush()
            return ids
      
      def DeleteByIds(self, ids):
            # ids = [1,]
            self.client.delete_entity_by_id(self.collection_name, ids) 
            self.Flush()

      def CreateIndex(self, index_type='IVF_FLAT', metric_type='L2', params={"dim": 4096}):
            ivf_param = {"index_type": index_type, "metric_type": metric_type, "params": params}
            self.client.create_index(self.collection_name, "Vec", ivf_param)
      
      def DropIndex(self, index_name):
            self.client.drop_index(self.collection_name, index_name)


      def Search(self, imgEncoding, collection_name, num=10, partition_tags=None, fields=None, timeout=None):
            # dsl = {
            #             "bool": {
            #                   "must":[
            #                         {
            #                         "term": {"A": [1, 2, 5]}
            #                         },
            #                         {
            #                         "range": {"B": {"GT": 1, "LT": 100}}
            #                         },
            #                         {
            #                         "vector": {
            #                               "Vec": {"topk": 10, "query": vectors[:1], "metric_type": "L2", "params": {"nprobe": 10}}
            #                         }
            #                         }
            #                   ]
            #             }
            #       }
            query_hybrid = {
                  "bool": {
                        "must": [
                              {
                              "vector": {
                                    "imgEncoding": {"topk": num, "query": [imgEncoding], "metric_type": "L2"}
                              }
                              }
                        ]
                  }
            }

            entities_list = self.client.search(collection_name, query_hybrid, fields=["imgID"])
            results = []
            for entities in entities_list:
                  for topk_film in entities:
                        result = {}
                        current_entity = topk_film.entity
                        # print("- id: {}".format(topk_film.id))
                        # print("- distance: {}".format(topk_film.distance))
                        # print("- imgID: {}".format(current_entity.imgID))
                        result["id"] = topk_film.id
                        result["distance"] = topk_film.distance
                        result["imgID"] = current_entity.imgID
                        results.append(result)
            return results

      def Flush(self):
            self.client.flush([self.collection_name])

      def GetStatus(self):
            info = self.client.get_collection_stats(self.collection_name)
            return info
